from transformers import *
import tensorflow as tf
import pandas as pd
import numpy as np
from tqdm import tqdm
import tensorflow_addons as tfa
import os
import json
import re


from encoder import BaseEncoder

from config_colab import Config

label_num = len(json.load(open(Config.label_dic_path)))


class F1ScoreNoZero(tfa.metrics.F1Score):



    def update_state(self, y_true, y_pred, sample_weight=None):

        y_true = tf.cast(y_true, tf.int32)
        y_true = tf.one_hot(y_true, label_num,axis=-1)

        y_true = y_true[:,:,1:]
        y_pred = y_pred[:,:,1:]


        if self.threshold is None:
            threshold = tf.reduce_max(y_pred, axis=-1, keepdims=True)
            y_pred = tf.logical_and(y_pred >= threshold, tf.abs(y_pred) > 1e-12)
        else:
            y_pred = y_pred > self.threshold

        y_true = tf.cast(y_true, self.dtype)
        y_pred = tf.cast(y_pred, self.dtype)

        def _weighted_sum(val, sample_weight):
            if sample_weight is not None:
                val = tf.math.multiply(val, tf.expand_dims(sample_weight, 1))
            return tf.reduce_sum(val, axis=self.axis)

        self.true_positives.assign_add(_weighted_sum(y_pred * y_true, sample_weight))
        self.false_positives.assign_add(
            _weighted_sum(y_pred * (1 - y_true), sample_weight)
        )
        self.false_negatives.assign_add(
            _weighted_sum((1 - y_pred) * y_true, sample_weight)
        )
        self.weights_intermediate.assign_add(_weighted_sum(y_true, sample_weight))


def build_model(encoder_list, label_num,lr):
    encoder_concat_list = []

    input_list_l = []


    for encoder in encoder_list:
        print(encoder)
        input_l,encoder_l = encoder.build()

        input_list_l.append(input_l)

        encoder_concat_list.append(encoder_l)

    x = tf.keras.layers.Concatenate()(encoder_concat_list)

    x = tf.keras.layers.BatchNormalization()(x)

    x = tf.keras.layers.Dense(label_num)(x)

    x = tf.keras.layers.Activation("softmax", name="output")(x)

    model = tf.keras.models.Model(input_list_l, x)

    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=lr),
                  loss=tf.keras.losses.SparseCategoricalCrossentropy(),
                  metrics=["acc", F1ScoreNoZero(num_classes=label_num, threshold=0.5, average="micro")])

    return model


